import os
import sys
import math
from bitarray import bitarray
from heapq import heappush, heappop

searchBufferSize = 4000
lookaheadBufferSize = 15


class heapNode(object):
    def __init__(self):
        self.left = None
        self.right = None
        self.counts = 0
        self.twelveBitInteger = ""

    def __lt__(self, other):
        return self.counts < other.counts


def printAndGetHuffmanTable(root, curPath, codes,
                            verbose):  # takes a huffman tree and generates huffman table and can also print it

    if root.left is None and root.right is None and root.twelveBitInteger:
        if verbose:
            print(curPath, "->", '{0:012b}'.format(root.twelveBitInteger))
        codes['{0:012b}'.format(root.twelveBitInteger)] = curPath
        return
    else:
        printAndGetHuffmanTable(root.left, curPath + "0", codes, verbose)
        printAndGetHuffmanTable(root.right, curPath + "1", codes, verbose)


def getBinaryFromTree(root, buffer):  # do preorder traversal to copy tree to file in binary
    if root.left is None and root.right is None:  # if leaf node
        buffer.append(False)
        bestDistanceLengthBinary = '{0:012b}'.format(root.twelveBitInteger)
        for bit in bestDistanceLengthBinary:
            if bit == '1':
                buffer.append(True)
            else:
                buffer.append(False)
    else:
        buffer.append(True)  # internal nodes will be 1
        getBinaryFromTree(root.left, buffer)
        getBinaryFromTree(root.right, buffer)


def getTreeFromBinary(buffer):  # do preorder traversal to construct new copy of binary tree
    nextBit = buffer.pop(0)
    if nextBit == False:  # if leaf node
        root = heapNode()
        binaryString = ""
        for i in range(12):  # read next twelve bytes to get full code
            bit = buffer.pop(0)
            if bit == True:
                binaryString += "1"
            else:
                binaryString += "0"

        root.twelveBitInteger = int(binaryString, 2)
    else:
        root = heapNode()
        root.left = getTreeFromBinary(buffer)
        root.right = getTreeFromBinary(buffer)

    return root


def getLongestSubstring(fileData, i):
    startSearch = max(0, i - searchBufferSize)
    endOfLookahead = min(i + lookaheadBufferSize, len(fileData) + 1)
    bestDistance = 0
    bestLength = 0

    mysearchbuffer = fileData[startSearch:i]
    lookaheadbuffer = fileData[i:endOfLookahead]
    lookaheadsubstringlength = 2
    while lookaheadsubstringlength < len(lookaheadbuffer) and len(mysearchbuffer) >= lookaheadsubstringlength:
        cursubstring = lookaheadbuffer[0:lookaheadsubstringlength]
        wassubstringfound = mysearchbuffer.rfind(cursubstring)
        if wassubstringfound != -1:
            lookaheadsubstringlength += 1
            if len(cursubstring) > bestLength:
                bestLength = len(cursubstring)
                distancefromendofsearchbuffer = (i - startSearch) - wassubstringfound
                bestDistance = distancefromendofsearchbuffer
                for j in range(len(cursubstring), len(lookaheadbuffer)):
                    if chr(lookaheadbuffer[j]) == chr(fileData[i - distancefromendofsearchbuffer + j]):
                        bestLength += 1
                    else:
                        break
        else:
            break

    if bestDistance > 0 and bestLength > 0:
        return (bestDistance, bestLength)
    else:
        return None


def compress(inputFile):
    fileData = None
    dataForTree = None
    i = 0
    j = 0
    distances = []
    outputBuffer = bitarray(endian='big')
    with open(inputFile, 'rb') as input_file:
        fileData = input_file.read()
        dataForTree = fileData

    while j < len(dataForTree):  #得到所有的距离数字，这样就可以用哈夫曼
        pointer2 = getLongestSubstring(dataForTree, j)
        if pointer2:
            (bestDistance, bestLength) = pointer2
            distances.append(bestDistance)
            j += bestLength
        else:
            j += 1

    counts = {}
    heap = []
    for distance in distances:  # 获取所有距离的计数
        if distance in counts:
            counts[distance] += 1
        else:
            counts[distance] = 1

    for key in counts:  #推入最小堆
        node = heapNode()
        node.twelveBitInteger = key
        node.counts = counts[key]
        heappush(heap, node)
    #通过弹出两个最小的数字来创建Huffman树，并创建一个新的内部节点，其计数是t的计数之和
    root = None
    while len(heap) > 1:
        smallestNode = heappop(heap)
        secondSmallestNode = heappop(heap)
        newRoot = heapNode()
        newRoot.counts = smallestNode.counts + secondSmallestNode.counts
        newRoot.twelveBitInteger = ''
        newRoot.left = smallestNode
        newRoot.right = secondSmallestNode
        root = newRoot
        heappush(heap, newRoot)
    codeDictionary = {}
    printAndGetHuffmanTable(root, '', codeDictionary, False)  # get the huffman table from tree
    treeBuffer = bitarray(endian='big')
    getBinaryFromTree(root, treeBuffer)  # convert tree into string of bits

    while i < len(
            fileData):  #不断移动搜索缓冲区和预查缓冲区的起始位置 [searchBuffer(max size = searchBufferSize)|i|lookaheadBuffer]
        pointer = getLongestSubstring(fileData, i)

        if pointer:
            # Format跟随1位标志，12位表示距离，4位表示子字符串的长度
            # 表示最大距离为4095，最大长度为15
            (bestDistance, bestLength) = pointer
            outputBuffer.append(True)  # 将1位标志连同距离和长度以二进制格式附加到
            huffmanCode = codeDictionary[
                '{0:012b}'.format(bestDistance)]  # get the huffman code corresponding to the 12 bit integer
            bestDistanceLengthBinary = huffmanCode + '{0:04b}'.format(bestLength)
            for bit in bestDistanceLengthBinary:
                if bit == '1':
                    outputBuffer.append(True)
                else:
                    outputBuffer.append(False)
            i += bestLength  #递增i无论i有多少字符
        else:
            #如果在缓冲区中没有找到匹配，则将标志和单个字符以二进制格式附加到outputBuffer中，后面的格式为1位作为标志，8位作为字符
            outputBuffer.append(False)
            outputBuffer.frombytes(bytes([fileData[i]]))
            i += 1
    # 确保newBuffer是8的倍数
    newBuffer = treeBuffer + outputBuffer  # 首先解析Treebuffer，然后解析其余的数据
    newBuffer.fill()

    # 将二进制fileData写入新文件
    inputFile = inputFile.replace('.', '')  # remove file extension from file name
    with open(inputFile + ".bin", 'wb') as outFile:
        outFile.write(newBuffer.tobytes())
        print("File compressed succesfully")

    # calculate compression ratio
    # originalFileSize = os.path.getsize(inputFile)
    # compressedFileSize = os.path.getsize(inputFile + ".bin")
    # compressionRatio = 1 - (compressedFileSize / originalFileSize)
    # print("Compression ratio: {:.2f}".format(compressionRatio))
    return None


def decompress(inputFile, outputFile):
    fileData = bitarray(endian='big')
    outputBuffer = []
    with open(inputFile, 'rb') as input_file:
        fileData.fromfile(input_file)
    codeDictionary = {}
    root = getTreeFromBinary(fileData)
    printAndGetHuffmanTable(root, "", codeDictionary, False)
    while len(fileData) >= 9:  # keep decompressing fileData while there are more then 8 bits left
        curbitsubstring = ""
        flag = fileData.pop(0)  # get the first bit of the current char/char sequence to be decoded
        if not flag:  # if just a single char then get next 8 bits and append to outputBuffer
            byte = fileData[0:8].tobytes()
            outputBuffer.append(byte)
            del fileData[0:8]
        else:  # next bits will be the huffman code

            bestLengthBinary = ""

            stop = False
            while stop == False:  # compare each bit subsequence and see if it is in the huffman table
                bit = fileData.pop(0)
                if bit == True:
                    curbitsubstring += "1"
                else:
                    curbitsubstring += "0"
                for key in codeDictionary:
                    if codeDictionary[key] == str(curbitsubstring):
                        twelvebitnumber = key  # if found in table then set the twelve bit number
                        stop = True

            for i in range(0, 4):  # get next 4 bits
                bit = fileData.pop(0)
                if bit == True:
                    bestLengthBinary += '1'
                else:
                    bestLengthBinary += '0'
            bestDistance = int(twelvebitnumber, 2)  # convert binary string to int
            bestLength = int(bestLengthBinary, 2)
            for i in range(bestLength):
                # get the next "length" of chars from the current Buffer that start at "distance" away and append to outputBuffer
                outputBuffer.append(outputBuffer[-bestDistance])

    # convert output outputBuffer to binary data to write to file
    decompressedfileData = b''.join(outputBuffer)

    # write binary fileData to new file
    with open(outputFile, 'wb') as outFile:
        outFile.write(decompressedfileData)
        print('File was decompressed successfully')
        return None


inputFile = 'GCA_014851415.1_ASM1485141v1_genomic.fna'



compress(inputFile)
decompress('GCA_014851415.1_ASM1485141v1_genomicfna.bin', 'new6666.txt')